#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
预处理（利用 2015–2022 的外生变量提高插值精度），
输出 2016–2022、严格对齐 employ_rate.csv 日期栅格的多元表：
  date, employment_rate, <外生变量...>

规则：
- 毕业生人数_万人：季度内常数（不做日内平滑）。缺失用相邻季度前向/后向填补。
- 其他变量：用 2015–2022 全窗口做日频插值，再在 2016–2022 的目标栅格取值。
- 如仍有缺失，用所在季度值兜底，确保无 NaN。
"""

import argparse, re, warnings, os
warnings.filterwarnings("ignore")

import pandas as pd
import numpy as np
from pandas.tseries.offsets import MonthEnd

PROCESS_YEARS_FULL = (2015, 2022)   # 外生变量插值用全窗口
OUTPUT_YEARS       = (2016, 2022)   # 仅输出 2016–2022

# ---------------- 时间解析（支持季度/年月/常见日期） ----------------
def _try_parse_quarter_str(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        zh_map = {"一":"1","二":"2","三":"3","四":"4","Ⅰ":"1","Ⅱ":"2","Ⅲ":"3","Ⅳ":"4"}
        for k,v in zh_map.items(): xs = xs.replace(k, v)
        m = re.match(r"^\s*(\d{4})\s*[-/ ]?\s*[Q第]?\s*([1-4])\s*[季度Q]?\s*$", xs)
        if m:
            y = int(m.group(1)); q = int(m.group(2))
            return pd.Timestamp(year=y, month=q*3, day=1) + MonthEnd(0)  # 季度末
        return pd.NaT
    return s.map(_to_dt)

def _try_parse_yearmonth(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        if re.match(r"^\d{6}$", xs):
            dt = pd.to_datetime(xs, format="%Y%m", errors="coerce")
            return dt + MonthEnd(0) if pd.notna(dt) else pd.NaT
        if re.match(r"^\d{4}[-/.]\d{1,2}$", xs):
            dt = pd.to_datetime(xs, errors="coerce")
            return (dt + MonthEnd(0)) if pd.notna(dt) else pd.NaT
        return pd.NaT
    return s.map(_to_dt)

def _coerce_datetime(df: pd.DataFrame, date_col: str) -> pd.DatetimeIndex:
    ser = df[date_col]
    dt = pd.to_datetime(ser, errors="coerce", infer_datetime_format=True)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_quarter_str(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_yearmonth(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    return pd.DatetimeIndex([pd.NaT]*len(ser))

def guess_and_parse_datetime(df: pd.DataFrame) -> pd.DatetimeIndex:
    for cand in ["date","时间","日期","ds","time"]:
        if cand in df.columns:
            dt = _coerce_datetime(df, cand)
            if dt.notna().any(): return dt
    low = {c.lower(): c for c in df.columns}
    if "year" in low and ("quarter" in low or "season" in low):
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        qn = low["quarter"] if "quarter" in low else low["season"]
        q = pd.to_numeric(df[qn], errors="coerce").clip(1,4)
        dt = pd.to_datetime(dict(year=y, month=(q*3).astype("Int64"), day=1),
                            errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    if "year" in low and "month" in low:
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        m = pd.to_numeric(df[low["month"]], errors="coerce").clip(1,12)
        dt = pd.to_datetime(dict(year=y, month=m, day=1),
                            errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    first = df.columns[0]
    dt = _coerce_datetime(df, first)
    if dt.notna().any(): return dt
    raise ValueError("无法识别时间列；可用列名：" + ", ".join(map(str, df.columns)))

# ---------------- 目标就业率（决定输出栅格与 y） ----------------
def ensure_percent_scale(s: pd.Series) -> pd.Series:
    s = pd.to_numeric(s, errors="coerce")
    if s.max(skipna=True) <= 1.0: s = s * 100.0
    return s.clip(lower=0, upper=100)

def load_target_grid_and_y(path: str):
    df = pd.read_csv(path)
    # 时间列
    date_col = None
    for c in ["date","时间","日期","ds","time"]:
        if c in df.columns: date_col = c; break
    if date_col is None:
        raise ValueError("employ_rate.csv 未找到时间列（date/时间/日期/ds/time）")
    # 目标列
    tcol = None
    for c in df.columns:
        if c in {"employment_rate","就业率","就业率（%）"}: tcol = c; break
    if tcol is None:
        for c in df.columns:
            cl = str(c).lower().replace("（%）","").replace("%","").strip()
            if "employment" in cl and "rate" in cl: tcol = c; break
    if tcol is None:
        raise ValueError("employ_rate.csv 未识别到目标列（employment_rate/就业率/就业率（%））")

    df = df.rename(columns={date_col:"date", tcol:"employment_rate"})
    df["date"] = pd.to_datetime(df["date"])
    # 仅输出 2016–2022，且保留原始栅格（不改采样点）
    df = df[(df["date"].dt.year>=OUTPUT_YEARS[0]) & (df["date"].dt.year<=OUTPUT_YEARS[1])]
    df = df.sort_values("date").drop_duplicates(subset=["date"])
    grid = pd.DatetimeIndex(df["date"])
    y = ensure_percent_scale(df["employment_rate"]).values
    return grid, pd.Series(y, index=grid, name="employment_rate")

# ---------------- 外生变量：对齐输出栅格（利用 2015–2022 提升插值精度） ----------------
def quarter_key(ts: pd.Timestamp) -> pd.Timestamp:
    q = (ts.month - 1) // 3 + 1
    return pd.Timestamp(year=ts.year, month=q * 3, day=1) + MonthEnd(0)

def load_allvars_on_grid(path: str, out_grid: pd.DatetimeIndex,
                         grad_col: str, wanted_cols=None) -> pd.DataFrame:
    """将外生变量映射到 out_grid（2016–2022），插值范围使用 2015–2022 的全窗口。"""
    raw = pd.read_csv(path)
    dt = guess_and_parse_datetime(raw)
    raw = raw.copy()
    raw["date"] = dt
    raw = raw.dropna(subset=["date"]).sort_values("date").set_index("date")

    # 仅用 2015–2022 这段作为插值与上下文
    raw = raw[(raw.index.year>=PROCESS_YEARS_FULL[0]) & (raw.index.year<=PROCESS_YEARS_FULL[1])]

    num = raw.select_dtypes(include=[np.number]).copy()
    if num.empty:
        raise ValueError("all_vars_yoy.csv 中没有数值列。")

    # 生成“全窗口日频插值”（给连续因子用）；毕业生人数列单独按季度处理
    daily_interp = num.resample("D").mean().interpolate("time").ffill().bfill()

    if wanted_cols is None:
        wanted_cols = list(num.columns)

    X = pd.DataFrame(index=out_grid)
    # 预生成季度映射
    q_map = pd.Series([quarter_key(d) for d in out_grid], index=out_grid)

    for col in wanted_cols:
        if col not in num.columns:
            raise ValueError(f"列 `{col}` 不在 all_vars_yoy.csv 的数值列中。")
        if grad_col and col == grad_col:
            # —— 毕业生人数：季度内常数 —— #
            s_q = num[col].copy().sort_index()
            s_q_filled = s_q.reindex(s_q.index.union(q_map.unique())).sort_index().ffill().bfill()
            X[col] = s_q_filled.reindex(q_map.values).values
        else:
            # —— 其他因子：先用全窗口做时间插值，再取出_grid；不足再用所在季度兜底 —— #
            s = pd.to_numeric(daily_interp[col], errors="coerce").reindex(out_grid)
            s = s.interpolate("time").ffill().bfill()
            if s.isna().any():
                s_q = num[col].copy().sort_index()
                s_q_filled = s_q.reindex(s_q.index.union(q_map.unique())).sort_index().ffill().bfill()
                s = s.fillna(pd.Series(s_q_filled.reindex(q_map.values).values, index=out_grid))
            X[col] = s.values

    # 最后兜底（理论上此时已无缺失）
    if X.isna().sum().sum() > 0:
        print("[WARN] 外生变量仍有缺失，已做 ffill/bfill 兜底：",
              X.isna().sum()[X.isna().sum()>0].to_dict())
        X = X.ffill().bfill()

    return X

# ---------------- 主流程 ----------------
def main():
    ap = argparse.ArgumentParser(description="预处理：用2015–2022提升插值精度，输出2016–2022对齐 employ_rate.csv 的多元表")
    ap.add_argument("--allvars", default="all_vars_yoy.csv", help="外生变量 CSV（含 CIER/毕业生人数_万人/GDP/三产占比等）")
    ap.add_argument("--target",  default="employ_rate.csv",  help="就业率 CSV（决定输出日期栅格）")
    ap.add_argument("--out",     default="pre_2016_2022_grid.csv", help="输出文件名")
    ap.add_argument("--grad-col", required=True, help="毕业生人数列名（季度内常数，不插值）")
    ap.add_argument("--cols", nargs="*", default=None, help="输出的外生变量列（顺序即输出顺序）。不填默认输出全部数值列。")
    args = ap.parse_args()

    if not os.path.exists(args.allvars) or not os.path.exists(args.target):
        raise SystemExit(f"找不到输入文件：{args.allvars} / {args.target}（可用 --allvars / --target 指定路径）")

    # 1) 读取目标就业率，获得输出栅格（2016–2022）
    out_grid, y_series = load_target_grid_and_y(args.target)

    # 2) 外生变量映射到输出栅格（插值用 2015–2022 窗口；毕业生人数季度常数）
    X = load_allvars_on_grid(args.allvars, out_grid, grad_col=args.grad_col, wanted_cols=args.cols)

    # 3) 组合输出
    out = pd.DataFrame({"date": out_grid, "employment_rate": y_series.values})
    # 列顺序：date, employment_rate, 外生变量（按 --cols 顺序）
    if args.cols:
        out = out.join(X[args.cols], on="date")
        # 再把剩余（未列在 --cols 中但存在的）附在最后，供后续分析
        rest = [c for c in X.columns if c not in args.cols]
        if rest:
            out = out.join(X[rest], on="date")
    else:
        out = out.join(X, on="date")

    # 确保数值列为 float
    for c in out.columns:
        if c not in {"date"}:
            out[c] = pd.to_numeric(out[c], errors="coerce")

    out.to_csv(args.out, index=False, encoding="utf-8-sig")
    print(f"[OK] 保存 {args.out}  rows={len(out)}  cols={out.shape[1]}")
    print("日期范围：", out['date'].min(), "~", out['date'].max())
    print("列名：", ", ".join(out.columns))

if __name__ == "__main__":
    main()
